import torch
from torchvision import transforms
from torchvision import datasets
import matplotlib.pyplot as plt
from PIL import Image

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((28, 28)),  # 调整照片大小
    transforms.Grayscale(num_output_channels=1),  # 转换为单通道灰度图
    transforms.ToTensor(),  # 转换为张量
])

# 加载模型
model = torch.load('模型文件.pth', weights_only=False)
model.eval()  # 将模型设置为评估模式
images = ['5.jpg', '4.png', '1.png', '6.png']
for image in images:
    # 自定义测试图片

    label = int(image[0])
    image = Image.open(image).convert('L')

    # ⬇️ 关键步骤：反转颜色（白变黑，黑变白）
    inverted_image = Image.eval(image, lambda x: 255 - x)

    # 应用 transform
    test_image = transform(inverted_image)

    test_image = test_image.unsqueeze(0)  # 添加一个批次维度

    output = model(test_image)

    # 获取预测结果
    _, predicted = torch.max(output, 1)
    predicted_class = predicted.item()
    print(f"预测{label}的类别是: {predicted_class}")

    # 显示图像
    image_display = test_image.squeeze(0).squeeze(0)  # 移除批次维度
    plt.imshow(image_display, cmap='gray')  # 使用灰度图显示
    plt.show()
